package com.fanxuankai.kafka.connect.transforms;

import org.apache.kafka.common.config.ConfigDef;
import org.apache.kafka.connect.connector.ConnectRecord;
import org.apache.kafka.connect.data.Field;
import org.apache.kafka.connect.data.Schema;
import org.apache.kafka.connect.data.SchemaBuilder;
import org.apache.kafka.connect.transforms.Transformation;
import org.apache.kafka.connect.transforms.util.SchemaUtil;

import java.util.Map;

/**
 * 抽取 key 为 Struct 类型
 *
 * @author fanxuankai
 */
public class ExtractStructKey<R extends ConnectRecord<R>> implements Transformation<R> {
    public static final String OVERVIEW_DOC = "If the key is not a struct type, it will be extracted as a struct type.";
    public static final ConfigDef CONFIG_DEF = new ConfigDef();

    @Override
    public void configure(Map<String, ?> props) {
    }

    @Override
    public R apply(R record) {
        if (record.keySchema() == null
                || record.keySchema().type() == Schema.Type.STRUCT) {
            return record;
        }
        SchemaBuilder updatedKeySchemaBuilder = SchemaUtil.copySchemaBasics(record.keySchema(), SchemaBuilder.struct());
        for (Field field : record.keySchema().fields()) {
            updatedKeySchemaBuilder.field(field.name(), field.schema());
        }
        Schema updatedKeySchema = updatedKeySchemaBuilder.build();
        return record.newRecord(record.topic(), record.kafkaPartition(),
                updatedKeySchema, record.key(), record.valueSchema(),
                record.value(), record.timestamp(), record.headers());
    }

    @Override
    public void close() {
    }

    @Override
    public ConfigDef config() {
        return CONFIG_DEF;
    }
}
